import re
import matplotlib.pyplot as plt

# 从日志文件中读取数据
log_file_path = "/home/ai999/project/div-align-dg/runs/thyroid_r101_c4_voc/log-div-align.txt"

# 初始化损失列表
total_loss = []
loss_classifier = []
loss_box_reg = []
loss_align_cls = []
loss_align_reg = []

# 正则表达式来提取损失信息
loss_pattern = re.compile(r"loss: ([\d.]+) .*loss_classifier: ([\d.]+) .*loss_box_reg: ([\d.]+) .*loss_align_cls: ([\d.]+) .*loss_align_reg: ([\d.]+)")

# 读取日志文件并提取损失信息
with open(log_file_path, "r") as log_file:
    for line in log_file:
        match = loss_pattern.search(line)
        if match:
            total_loss.append(float(match.group(1)))
            loss_classifier.append(float(match.group(2)))
            loss_box_reg.append(float(match.group(3)))
            loss_align_cls.append(float(match.group(4)))
            loss_align_reg.append(float(match.group(5)))

# 去除前 10 个迭代的异常高损失值
cut_off = 10  # 去掉前 10 次迭代的数据
total_loss = total_loss[cut_off:]
loss_classifier = loss_classifier[cut_off:]
loss_box_reg = loss_box_reg[cut_off:]
loss_align_cls = loss_align_cls[cut_off:]
loss_align_reg = loss_align_reg[cut_off:]

# 创建绘图
plt.figure(figsize=(15, 10))

# 绘制总损失曲线
plt.subplot(2, 3, 1)
plt.plot(total_loss, label="Total Loss", color='blue')
plt.title('Total Loss')
plt.xlabel('Iterations')
plt.ylabel('Loss')
plt.legend()

# 绘制分类损失曲线
plt.subplot(2, 3, 2)
plt.plot(loss_classifier, label="Classification Loss", color='green')
plt.title('Classification Loss')
plt.xlabel('Iterations')
plt.ylabel('Loss')
plt.legend()

# 绘制边界框回归损失曲线
plt.subplot(2, 3, 3)
plt.plot(loss_box_reg, label="Box Regression Loss", color='red')
plt.title('Box Regression Loss')
plt.xlabel('Iterations')
plt.ylabel('Loss')
plt.legend()

# 绘制对齐分类损失曲线
plt.subplot(2, 3, 4)
plt.plot(loss_align_cls, label="Alignment Classification Loss", color='purple')
plt.title('Alignment Classification Loss')
plt.xlabel('Iterations')
plt.ylabel('Loss')
plt.legend()

# 绘制对齐回归损失曲线
plt.subplot(2, 3, 5)
plt.plot(loss_align_reg, label="Alignment Regression Loss", color='orange')
plt.title('Alignment Regression Loss')
plt.xlabel('Iterations')
plt.ylabel('Loss')
plt.legend()

# 调整布局并保存图片
plt.tight_layout()
plt.savefig('div-align-loss.png')
